# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import collections
import os
import datetime
import logging
import sys
import numpy
import pandas

from absl import app, flags
from concurrent.futures import ProcessPoolExecutor

from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str
from coin.support.proto_log.logic.util import run_from_og_log_archive
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog
import google.protobuf.json_format as json_format
from coin.strategy.mm.tool.spot_pnl_setting import SpotPnlSetting

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


class SpotPnlCalculator(object):
  def __init__(self, trading_date, strat_setting):
    self._trading_date = trading_date
    self._market_type = strat_setting.market_type
    self._exchange = strat_setting.exchange
    self._machine = strat_setting.machine
    self._maxpos_config = strat_setting.max_pos_map
    self._symbol2prod = {}
    self._quote_currency = None

    self._hold_map = {}
    self._avail_map = {}
    self._holds = {}
    self._avails = {}

    self._last_n_balance_maps = collections.deque()
    self._balance_map = {}
    self._base2quote = {}
    self._balances = {}
    self._base2quotes = {}
    self._maxpos_map = {}
    self._maxposs = {}
    self._timestamp_midnight = int(self._trading_date.timestamp() * 1e9)
    self._next_timestamp = self._timestamp_midnight
    self._next_timestamp_map = {}
    self._processed_hour = {}
    self._prev_balance_map = {}
    self._prev_price_map = {}

    self._max_dtstr = None
    self._holding = None
    self._prev_holding = None

    next_date = trading_date + datetime.timedelta(days=1)
    curtime = trading_date
    while curtime <= next_date:
      next_time = curtime + datetime.timedelta(seconds=1)
      self._next_timestamp_map[curtime.strftime("%H%M%S")] = int(next_time.timestamp() * 1e9)
      curtime = next_time

  def on_log2(self, timestamp, log):
    pb = OrderGatewayLog()
    try:
      pb.ParseFromString(log)
    except Exception:
      print(log)
      return
    if pb.type in [OrderGatewayLog.BALANCE, OrderGatewayLog.ORDER_EVENT]:
      # this comes too frequently... skips for the efficiency.
      if pb.type == OrderGatewayLog.BALANCE:
        if timestamp < self._next_timestamp:
          return
      log_json = json_format.MessageToDict(pb, preserving_proto_field_name=True)
      self.on_log_json(timestamp, log_json)

  def update_product_map(self, symbol):
    if symbol not in self._symbol2prod:
      if self._market_type == 'Futures' and self._exchange == 'Okex':
        product = OkexFuturesProduct.FromStr(symbol, self._trading_date)
      else:
        product = generate_product_from_str(self._exchange, None, symbol)
      self._symbol2prod[symbol] = product
      if self._quote_currency is None:
        self._quote_currency = product.quote.currency
        self._base2quote[self._quote_currency] = 1
      else:
        assert self._quote_currency == product.quote.currency, (self._quote_currency,
                                                                product.quote.currency)

  def add_balance_row(self, timestamp):
    hourstr = pandas.DatetimeIndex([timestamp]).strftime("%H")[0]
    if hourstr not in self._processed_hour:
      if len(self._processed_hour) > 0:
        self.export_log(preexport=True)
      self._processed_hour[hourstr] = True
      logging.info("processing %s %s:00", self._trading_date, hourstr)
    dtstr = pandas.DatetimeIndex([timestamp]).strftime("%H%M%S")[0]
    if self._max_dtstr is None or dtstr > self._max_dtstr:
      self._max_dtstr = dtstr
      # removing freaking spikes on the balances...
      # for example, if you buy BTC with USDT,
      # USDT can decrease first, and then BTC increases,
      # which can cause race condition.
      balance_tuple, _ = collections.Counter(
          [tuple(bal_map.items()) for bal_map in self._last_n_balance_maps]).most_common()[0]
      self._balance_map = dict(balance_tuple)
      self._next_timestamp = self._next_timestamp_map[dtstr]
      self._balances[dtstr] = self._balance_map.copy()
      self._holds[dtstr] = self._hold_map.copy()
      self._avails[dtstr] = self._avail_map.copy()
      self._base2quotes[dtstr] = self._base2quote.copy()
      self._maxposs[dtstr] = self._maxpos_map.copy()
      self._prev_balance_map = self._balance_map.copy()
      self._prev_holding = self._holding
      self._prev_price_map = self._base2quote.copy()
      self._last_n_balance_maps.clear()

  def on_log_json(self, timestamp, log_json):
    if log_json['type'] == 'BALANCE':
      balance = log_json['balance']
      for each_balance in balance['each_balance']:
        self._balance_map[each_balance['currency']] = round(each_balance['total'], 8)
        self._hold_map[each_balance['currency']] = round(each_balance['hold'], 8)
        self._avail_map[each_balance['currency']] = round(each_balance['available'], 8)
        self._maxpos_map[each_balance['currency']] = self._maxpos_config.get(
            each_balance['currency'], None)
      self._last_n_balance_maps.append(self._balance_map.copy())
      self.add_balance_row(timestamp)
    elif log_json['type'] == 'ORDER_EVENT':
      event = log_json['event']
      if event.get('order_qty', 0) > 0:
        self.update_product_map(event['symbol'])
        product = self._symbol2prod[event['symbol']]
        self._base2quote[product.base.currency] = event['order_price']

  def export_log(self, preexport=False):
    baldf = pandas.DataFrame(self._balances).T
    holddf = pandas.DataFrame(self._holds).T
    availdf = pandas.DataFrame(self._avails).T
    prcdf = pandas.DataFrame(self._base2quotes).T
    maxbaldf = pandas.DataFrame(self._maxposs).T

    sel = (numpy.isnan(prcdf).sum(1) == 0)
    port_value = (prcdf * baldf).sum(1)[sel]
    reserve_value = (prcdf * maxbaldf).sum(1)[sel] * 0.5
    smooth_window = 20
    port_value = port_value.rolling(window=smooth_window).median()[smooth_window:]
    reserve_value = reserve_value[smooth_window:]
    neut_value = port_value - reserve_value
    accum_pnl = neut_value - neut_value.iloc[0]

    if not preexport:
      logging.info("[time] %s, %s->%s" %
                   (self._trading_date, accum_pnl.index.tolist()[0], accum_pnl.index.tolist()[-1]))
      logging.info("[pnl(%s)]: %s -> %s", 0, self._quote_currency, accum_pnl.iloc[-1])
      logging.info("[balance(%s)]: %s -> %s(%s)",
                   self._quote_currency,
                   port_value.iloc[0],
                   port_value.iloc[-1],
                   port_value.iloc[-1] - port_value.iloc[0])
      logging.info("[reserve(%s)]: %s -> %s(%s)",
                   self._quote_currency,
                   reserve_value.iloc[0],
                   reserve_value.iloc[-1],
                   reserve_value.iloc[-1] - reserve_value.iloc[0])

    dti = pandas.to_datetime(self._trading_date.strftime("%Y%m%d")
                             + pandas.Series(accum_pnl.index).astype('str'),
                             format="%Y%m%d%H%M%S")
    os.makedirs("pnl_stat", exist_ok=True)
    plt.rcParams['font.size'] = 6
    plt.subplot(311)
    plt.title("%s %s pnl(%s)" % (self._trading_date, FLAGS.strategy_name, self._quote_currency))
    plt.plot(dti, accum_pnl.tolist(), linewidth=0.5)
    plt.xticks(rotation=20)
    plt.subplot(312)
    plt.title("balance(%s)" % self._quote_currency)
    plt.plot(dti, port_value.tolist(), linewidth=0.5)
    plt.xticks(rotation=20)
    plt.subplot(313)
    plt.title("reserve(%s)" % self._quote_currency)
    plt.plot(dti, reserve_value.tolist(), linewidth=0.5)
    plt.xticks(rotation=20)
    plt.tight_layout()
    plt.savefig("pnl_stat/%s_%s.png" % (FLAGS.strategy_name, self._trading_date), dpi=150)
    plt.close()


def launch(trading_date, root_dir, strategy_name):
  start_time = trading_date  # UTC 0 hour
  # 24 will omit some logs..
  hours = 25

  strat_setting = SpotPnlSetting.from_name(strategy_name, start_time)

  strategy = SpotPnlCalculator(trading_date, strat_setting)
  logging.info("start log analysis")
  run_from_og_log_archive(strategy.on_log2,
                          strat_setting.market_type,
                          strat_setting.exchange,
                          start_time,
                          start_time + datetime.timedelta(hours=hours),
                          root_dir,
                          strat_setting.machine,
                          owner=None)
  strategy.export_log()
  return 0


def get_trading_dates(trading_date_str):
  if trading_date_str.find("-") >= 0:
    tdfrom, tdto = trading_date_str.split("-")
    return pandas.date_range(tdfrom, tdto).to_pydatetime()
  else:
    return [datetime.datetime.strptime(trading_date_str, '%Y%m%d')]


def main(_):
  # import warnings; warnings.filterwarnings("error")
  trading_date = FLAGS.trading_date
  strategy_name = FLAGS.strategy_name
  root_dir = os.path.expanduser(FLAGS.root_dir)
  print('Running for %s %s ...' % (trading_date, strategy_name))
  sys.stdout.flush()
  pandas.options.display.width = 200
  trading_dates = get_trading_dates(trading_date)
  if len(trading_dates) == 1:
    launch(trading_dates[0], root_dir, strategy_name)
  else:
    with ProcessPoolExecutor(max_workers=4) as executor:
      futures = []
      for trading_date in trading_dates:
        futures.append(executor.submit(launch, trading_date, root_dir, strategy_name))
      for future in futures:
        future.result()


if __name__ == '__main__':
  flags.DEFINE_string('exchange', 'Okex', 'Exchange name.')

  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('strategy_name', None, 'strategy name')

  flags.DEFINE_string('root_dir', '~/data/og_proto_log', 'root_dir.')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
